package org.iscas.tcse.util;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.ml.linalg.DenseVector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class CsvToVector implements MapFunction<String, DenseVector> {
    private static final Logger LOG = LoggerFactory.getLogger(CsvToVector.class);
    private final String delimiter;
    private final int numOfFeatures;

    private final boolean labelFirst;

    public CsvToVector(String delimiter, int numOfFeatures, boolean labelFirst) {
        this.delimiter = delimiter;
        this.numOfFeatures = numOfFeatures;
        this.labelFirst = labelFirst;
        LOG.info("delimiter is {}, numOfFeatures is {}", delimiter, numOfFeatures);
    }

    @Override
    public DenseVector map(String value) throws Exception {
        String[] split = value.split(this.delimiter);

        double[] features = new double[numOfFeatures + 1];

        if (labelFirst) {
            features[numOfFeatures] = Double.parseDouble(split[0]);
        } else {
            features[numOfFeatures] = Double.parseDouble(split[numOfFeatures]);
        }
        for (int i = 0; i < numOfFeatures; i++) {
            if (labelFirst) {
                features[i] = Double.parseDouble(split[i + 1]);
            } else {
                features[i] = Double.parseDouble(split[i]);
            }
        }

        return new DenseVector(features);
    }
}
